from abc import ABC, abstractmethod
from typing import List, Dict, Any, Optional, Union
import pandas as pd
import numpy as np

from qtorch.core.strategy import Strategy
from qtorch.core.interfaces import IStrategy


class QStrategy(Strategy, IStrategy):
    """
    简化的策略基类，提供更直观的策略开发接口
    
    包含三个主要方法：
    1. initialize: 在回测开始前调用，用于初始化策略
    2. next: 每个数据点调用一次，用于生成交易信号
    3. generate_signals: 实现Strategy接口，内部调用next方法
    """
    
    def __init__(self):
        """初始化策略"""
        self.position = 0  # 当前持仓状态
        self.prices = []   # 历史价格
        self.data = None   # 当前数据点
        self.initialized = False
        self._indicators = {}  # 指标缓存
        
    def initialize(self):
        """
        策略初始化，在回测开始前调用
        用于初始化指标、参数等
        """
        pass
        
    @abstractmethod
    def next(self) -> int:
        """
        核心策略逻辑，处理每个数据点
        
        返回:
            int: 交易信号 1(买入)、0(持有)、-1(卖出)
        """
        pass
        
    def generate_signals(self, data: pd.DataFrame) -> np.ndarray:
        """
        实现Strategy接口，批量生成信号
        
        参数:
            data: 市场数据DataFrame
            
        返回:
            np.ndarray: 交易信号数组
        """
        # 确保初始化
        if not self.initialized:
            self.initialize()
            self.initialized = True
            
        signals = []
        
        try:
            # 转换为DataFrame
            if not isinstance(data, pd.DataFrame):
                df = pd.DataFrame(data)
            else:
                df = data
                
            # 确保包含必要的列
            required_columns = ['close']
            for col in required_columns:
                if col not in df.columns:
                    raise ValueError(f"输入数据必须包含{col}列")
            
            # 逐行处理数据
            for i in range(len(df)):
                # 更新当前数据点
                self.data = df.iloc[i]
                
                # 更新价格历史
                if 'close' in self.data:
                    self.prices.append(self.data['close'])
                
                # 生成信号
                signal = self.next()
                signals.append(signal)
                
                # 更新持仓状态
                if signal == 1:  # 买入
                    self.position = 1
                elif signal == -1:  # 卖出
                    self.position = -1
                # 持有不变
            
            return np.array(signals)
            
        except Exception as e:
            import logging
            logging.error(f"策略信号生成失败: {str(e)}", exc_info=True)
            return np.zeros(len(data))
    
    # 实用指标计算方法
    def SMA(self, period: int, price_field: str = 'close') -> float:
        """
        简单移动平均线
        
        参数:
            period: 周期
            price_field: 价格字段
            
        返回:
            float: 当前SMA值
        """
        key = f"SMA_{period}_{price_field}"
        
        if len(self.prices) < period:
            return None
            
        if key not in self._indicators:
            self._indicators[key] = []
            
        current_value = sum(self.prices[-period:]) / period
        self._indicators[key].append(current_value)
        
        return current_value
        
    def EMA(self, period: int, price_field: str = 'close') -> float:
        """
        指数移动平均线
        
        参数:
            period: 周期
            price_field: 价格字段
            
        返回:
            float: 当前EMA值
        """
        key = f"EMA_{period}_{price_field}"
        
        if len(self.prices) < period:
            return None
            
        if key not in self._indicators:
            # 初始化EMA为简单平均
            self._indicators[key] = [sum(self.prices[:period]) / period]
            
        # 计算当前EMA
        alpha = 2 / (period + 1)
        prev_ema = self._indicators[key][-1]
        current_ema = alpha * self.prices[-1] + (1 - alpha) * prev_ema
        self._indicators[key].append(current_ema)
        
        return current_ema
        
    def RSI(self, period: int = 14) -> float:
        """
        相对强弱指标
        
        参数:
            period: 周期
            
        返回:
            float: 当前RSI值
        """
        key = f"RSI_{period}"
        
        if len(self.prices) <= period:
            return None
            
        # 计算涨跌幅
        deltas = np.diff(self.prices)
        
        # 分离上涨和下跌
        gains = deltas.copy()
        losses = deltas.copy()
        gains[gains < 0] = 0
        losses[losses > 0] = 0
        losses = abs(losses)
        
        # 如果数据长度足够
        if len(gains) >= period:
            avg_gain = gains[-period:].mean()
            avg_loss = losses[-period:].mean()
            
            if avg_loss == 0:
                return 100
                
            rs = avg_gain / avg_loss
            rsi = 100 - (100 / (1 + rs))
            
            return rsi
            
        return None
        
    def MACD(self, fast_period: int = 12, slow_period: int = 26, signal_period: int = 9) -> Dict[str, float]:
        """
        MACD指标
        
        参数:
            fast_period: 快线周期
            slow_period: 慢线周期
            signal_period: 信号线周期
            
        返回:
            Dict: 包含macd, signal, histogram的字典
        """
        # 计算快线和慢线
        fast_ema = self.EMA(fast_period)
        slow_ema = self.EMA(slow_period)
        
        if fast_ema is None or slow_ema is None:
            return None
            
        # 计算MACD线
        key = f"MACD_{fast_period}_{slow_period}_{signal_period}"
        
        if key not in self._indicators:
            self._indicators[key] = {
                'macd_line': [],
                'signal_line': [],
                'histogram': []
            }
            
        # 计算MACD值
        macd_line = fast_ema - slow_ema
        self._indicators[key]['macd_line'].append(macd_line)
        
        # 计算信号线 (MACD的EMA)
        if len(self._indicators[key]['macd_line']) >= signal_period:
            macd_data = self._indicators[key]['macd_line'][-signal_period:]
            signal_line = sum(macd_data) / signal_period
            self._indicators[key]['signal_line'].append(signal_line)
            
            # 计算MACD柱状图
            histogram = macd_line - signal_line
            self._indicators[key]['histogram'].append(histogram)
            
            return {
                'macd': macd_line,
                'signal': signal_line,
                'histogram': histogram
            }
            
        return None